from IPy import IP
from flask import Flask, request
from flask_caching import Cache

from app.common.utils import get_remote_address
from app.models import IpFilter
from app.common.enums import IpFilterType
from app.error import IpFilterError


class AppIpFilter:
    """全局IP拦截
    """

    def __init__(self, app: Flask | None = None, cache: Cache | None = None):
        if app is not None:
            self.init_app(app, cache)

    def init_app(self, app: Flask, cache: Cache) -> None:
        """全局IP拦截初始化"""
        # 设置默认参数
        app.config.setdefault('IP_FILTER_ENABLED', True)

        @app.before_request
        def limit_remote_addr():
            if app.config.get('IP_FILTER_ENABLED'):
                ip = get_remote_address(request)
                app.logger.debug(f'==> before_request,ip={ip}')
                # 查询过滤的ip列表
                for ip_obj in get_ip_filter():
                    if ip in IP(ip_obj.network_segment):
                        # 白名单，触发就直接结束校验，通过
                        if ip_obj.filter_type == IpFilterType.white.value:
                            app.logger.warn(f'{ip}触发白名单网段{ip_obj.network_segment}')
                            break
                        else:
                            # 黑名单，触发就直接返回错误消息
                            app.logger.warn(f'{ip}触发黑名单网段{ip_obj.network_segment}')
                            raise IpFilterError()

        @cache.memoize(timeout=60)
        def get_ip_filter():
            """查询过滤的ip列表，带缓存
            """
            return IpFilter.query.filter_by(is_enabled=True).all()
